#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep
import grpc
from datetime import datetime

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

# Topo Data Structures Renamed
ENDPOINT_PORTS = {
    's1': 1,
    's2': 1,
    's3': 1
}

INTER_DEVICE_PORTS = {
    's1': {
        's2': 2,
        's3': 3
    },
    's2': {
        's1': 2,
        's3': 3
    },
    's3': {
        's1': 2,
        's2': 3
    }
}

NETWORK_ENDPOINTS = {
    'e1': {'ip': '10.0.1.1', 'mac': '08:00:00:00:01:11', 'device': 's1'},
    'e2': {'ip': '10.0.2.2', 'mac': '08:00:00:00:02:22', 'device': 's2'},
    'e3': {'ip': '10.0.3.3', 'mac': '08:00:00:00:03:33', 'device': 's3'}
}

OVERLAY_IDS = {
    ('s1', 's2'): 100,
    ('s2', 's1'): 101,
    ('s1', 's3'): 200,
    ('s3', 's1'): 201,
    ('s2', 's3'): 300,
    ('s3', 's2'): 301
}

def provisionOverlayPath(info_map, ingress_dev, egress_dev, path_id, dst_mac_addr, dst_ip_addr):
    try:
        print(f"Provisioning overlay path from {ingress_dev.name} to {egress_dev.name} with ID {path_id}")

        # 1. Ingress Device: Install Path Initiation Rule (Encapsulation)
        table_record = info_map.buildTableEntry(
            table_name="IngressProcessing.IP_Destination_Lookup",
            match_fields={"pkt_hdr.ip.destination_ip": (dst_ip_addr, 32)},
            action_name="IngressProcessing.Prepare_Encapsulation",
            action_params={"dest_tunnel_id": path_id})

        ingress_dev.WriteTableEntry(table_record)
        print(f"  Installed initiation rule on {ingress_dev.name}")

        # 2. Ingress/Transit Device: Install Transit Forwarding Rule
        if ingress_dev.name != egress_dev.name:
            outbound_port = INTER_DEVICE_PORTS[ingress_dev.name][egress_dev.name]
            table_record = info_map.buildTableEntry(
                table_name="IngressProcessing.Overlay_ID_Match",
                match_fields={"pkt_hdr.tunnel.tunnel_id": path_id},
                action_name="IngressProcessing.Tunnel_Pass_Through",
                action_params={"out_port": outbound_port})

            ingress_dev.WriteTableEntry(table_record)
            print(f"  Installed transit rule on {ingress_dev.name}")

        # 3. Egress Device: Install Path Termination Rule (Decapsulation)
        table_record = info_map.buildTableEntry(
            table_name="IngressProcessing.Overlay_ID_Match",
            match_fields={"pkt_hdr.tunnel.tunnel_id": path_id},
            action_name="IngressProcessing.Setup_Decapsulation",
            action_params={
                "final_dest_mac": dst_mac_addr,
                "out_port": ENDPOINT_PORTS[egress_dev.name]
            })

        egress_dev.WriteTableEntry(table_record)
        print(f"  Installed termination rule on {egress_dev.name}")

        return True

    except Exception as e:
        print(f"  Error provisioning path rules: {e}")
        return False

def configureDirectConnect(info_map, devices):
    print("\n----- Configuring direct connection rules -----")

    device_map = {dev.name: dev for dev in devices}

    for endpoint_name, endpoint_data in NETWORK_ENDPOINTS.items():
        device_name = endpoint_data['device']
        if device_name in device_map:
            table_record = info_map.buildTableEntry(
                table_name="IngressProcessing.IP_Destination_Lookup",
                match_fields={"pkt_hdr.ip.destination_ip": (endpoint_data['ip'], 32)},
                action_name="IngressProcessing.Process_L3_Forwarding",
                action_params={
                    "next_hop_mac": endpoint_data['mac'],
                    "out_port": ENDPOINT_PORTS[device_name]
                })

            device_map[device_name].WriteTableEntry(table_record)
            print(f"Configured direct connection rule on {device_name} for {endpoint_data['ip']}")

def retrieveCounterData(info_map, dev, counter_id, lookup_index):
    try:
        counter_name = "Encapsulation_Counter" if counter_id == "Encapsulation_Counter" else "Decapsulation_Counter"
        p4_counter_id = info_map.get_counters_id(counter_name)
        
        for response in dev.ReadCounters(p4_counter_id, lookup_index):
            for entity in response.entities:
                counter = entity.counter_entry
                if counter.data:
                    return counter.data.packet_count, counter.data.byte_count
        return 0, 0
    except Exception as e:
        print(f"Error retrieving counter {counter_name}[{lookup_index}] on {dev.name}: {e}")
        return 0, 0

def logTunnelFlow(flow_tag, flow_stats):
    os.makedirs("flow_logs", exist_ok=True)
    filename = f"flow_logs/{flow_tag}.log"

    report_lines = []

    for flow_direction, flow_data in flow_stats.items():
        if flow_data['packets'] > 0:
            if flow_direction == "encap":
                report_lines.append(f"{flow_data['dst_dev']} received {flow_data['packets']} packets, flow ID: {flow_data['path_id']}")
            else:
                report_lines.append(f"{flow_data['src_dev']} sent {flow_data['packets']} packets, flow ID: {flow_data['path_id']}")

    if report_lines:
        with open(filename, "w") as f:
            for line in report_lines:
                f.write(line + "\n")

def checkAllFlowStats(info_map, devices):
    print(f'\n----- Retrieving flow statistics -----')

    device_map = {dev.name: dev for dev in devices}

    path_tags = {
        ('s1', 's2'): 'S1_to_S2',
        ('s2', 's1'): 'S2_to_S1',
        ('s1', 's3'): 'S1_to_S3',
        ('s3', 's1'): 'S3_to_S1',
        ('s2', 's3'): 'S2_to_S3',
        ('s3', 's2'): 'S3_to_S2'
    }

    for (src_dev, dst_dev), flow_tag in path_tags.items():
        if src_dev in device_map and dst_dev in device_map:
            path_id = OVERLAY_IDS[(src_dev, dst_dev)]

            encap_pkts, encap_bytes = retrieveCounterData(info_map, device_map[src_dev],
                                                          "Encapsulation_Counter", path_id)
            decap_pkts, decap_bytes = retrieveCounterData(info_map, device_map[dst_dev],
                                                          "Decapsulation_Counter", path_id)

            flow_stats = {
                'encap': {
                    'src_dev': src_dev,
                    'dst_dev': dst_dev,
                    'path_id': path_id,
                    'packets': encap_pkts,
                    'bytes': encap_bytes
                },
                'decap': {
                    'src_dev': src_dev,
                    'dst_dev': dst_dev,
                    'path_id': path_id,
                    'packets': decap_pkts,
                    'bytes': decap_bytes
                }
            }

            if encap_pkts > 0 or decap_pkts > 0:
                print(f"--- {src_dev} -> {dst_dev} ({flow_tag}) ---")
                if encap_pkts > 0:
                    print(f"{src_dev} Encapsulation_Counter {path_id}: {encap_pkts} pkts")
                if decap_pkts > 0:
                    print(f"{dst_dev} Decapsulation_Counter {path_id}: {decap_pkts} pkts")

                logTunnelFlow(flow_tag, flow_stats)

def setupAllOverlayPaths(info_map, devices):
    print("\n----- Setting up all overlay paths -----")

    device_map = {dev.name: dev for dev in devices}

    successful_paths = 0
    for (src_dev, dst_dev), path_id in OVERLAY_IDS.items():
        if src_dev in device_map and dst_dev in device_map:
            dest_endpoint = None
            for endpoint_name, endpoint_data in NETWORK_ENDPOINTS.items():
                if endpoint_data['device'] == dst_dev:
                    dest_endpoint = endpoint_name
                    break

            if dest_endpoint:
                if provisionOverlayPath(
                    info_map,
                    device_map[src_dev],
                    device_map[dst_dev],
                    path_id,
                    NETWORK_ENDPOINTS[dest_endpoint]['mac'],
                    NETWORK_ENDPOINTS[dest_endpoint]['ip']
                ):
                    successful_paths += 1

    print(f"Successfully configured {successful_paths} overlay paths")

def main_orchestrator(p4info_path, json_path):
    info_map = p4runtime_lib.helper.P4InfoHelper(p4info_path)

    try:
        # Establish Device Connections
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1', address='127.0.0.1:50051', device_id=0,
            proto_dump_file='flow_logs/s1_proto_msgs.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2', address='127.0.0.1:50052', device_id=1,
            proto_dump_file='flow_logs/s2_proto_msgs.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3', address='127.0.0.1:50053', device_id=2,
            proto_dump_file='flow_logs/s3_proto_msgs.txt')

        devices = [s1, s2, s3]

        # Initialize and Configure Devices
        for dev in devices:
            dev.MasterArbitrationUpdate()

        for dev in devices:
            dev.SetForwardingPipelineConfig(p4info=info_map.p4info,
                                            bmv2_json_file_path=json_path)

        # Install Data Plane Rules
        configureDirectConnect(info_map, devices)
        setupAllOverlayPaths(info_map, devices)

        # Wait for network stabilization
        sleep(5)

        # Start Flow Monitoring Loop
        print('----- Starting continuous flow monitoring -----')
        while True:
            sleep(2)
            checkAllFlowStats(info_map, devices)

    except KeyboardInterrupt:
        print("\nOrchestrator shutting down.")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Network Flow Orchestrator')
    parser.add_argument('--p4info', help='P4 Info file path', default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--json', help='BMv2 JSON file path', default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        print(f"Required P4 Info file not found: {args.p4info}")
        exit(1)
    if not os.path.exists(args.json):
        print(f"Required JSON file not found: {args.json}")
        exit(1)

    os.makedirs("flow_logs", exist_ok=True)
    main_orchestrator(args.p4info, args.json)